import websocket
import json
import time
import numpy as np
import cv2 as cv
from PIL import Image, ImagePalette

try:
    import thread
except ImportError:
    import _thread as thread

colors = [
    (0xf5, 0xde, 0xb3),  # 未走过的路
    (0x8a, 0x2b, 0xe2),  # 墙壁
    (0xd3, 0xdb, 0x64),  # 当前位置
    (0x12, 0x8b, 0x37),  # 走过
    (0x32, 0x12, 0x8b),  # 重复经过
    (0xff, 0x00, 0x00),
    (0xff, 0x00, 0x00),
    (0xff, 0x00, 0x00),
    (0xff, 0x00, 0x00),
    (0xc4, 0xc4, 0xc4),  # 迷雾
    (0xff, 0x00, 0x00),
    (0x00, 0xff, 0xff),
    (0xff, 0xff, 0x00),
]


def on_message(ws, message):

    """服务器有数据更新时，主动推送过来的数据"""
    print(message)
    msg = json.loads(message)

    def run(*args):
        if msg.get("action") == "subscribe-accept-handler":
            request_next_frame = {
                "action": "request-next-frame-handler",
                "recordId": msg.get("recordId")
            }

            ws.send(json.dumps(request_next_frame))  # 发送数据(必须为str类型)

        if msg.get("action") == "response-next-frame-handler":
            maze = np.array(msg.get("maze"), dtype='int8')
            im = Image.fromarray(maze)
            im = im.convert("P")  # P代表索引图片
            palette = np.array(colors).reshape(-1).tolist()
            print(palette)
            im.putpalette(palette)
            scale = 70
            resized_img = im.resize((im.width * scale, im.height * scale), Image.BICUBIC)
            resized_img = resized_img.convert("RGB")
            frame = cv.cvtColor(np.asarray(resized_img), cv.COLOR_RGB2BGR)
            cv.imwrite('/home/xlj/received-frame.png', frame)

            request_next_frame = {
                "action": "request-next-frame-handler",
                "recordId": msg.get("recordId")
            }

            ws.send(json.dumps(request_next_frame))  # 发送数据(必须为str类型)

        time.sleep(1)

    thread.start_new_thread(run, ())


def on_error(ws, error):

    """程序报错时，就会触发on_error事件"""

    print(error)


def on_close(ws):

    """关闭websocket连接后，打印此消息"""

    print("### closed ###")


def on_open(ws):

    """连接到服务器之后就会触发on_open事件，这里用于send数据"""

    def run(*args):

        subscribe = {
            "action": "challenge-live-request-handler",
            "recordId": 123
        }

        ws.send(json.dumps(subscribe))  # 发送数据(必须为str类型)

        time.sleep(1)
        #
        # ws.close()  # 关闭连接

    thread.start_new_thread(run, ())


def run_1():

    websocket.enableTrace(False)

    ws = websocket.WebSocketApp("ws://localhost:9000/maze-endpoint/python",
                                on_message=on_message,
                                on_error=on_error,
                                on_close=on_close)

    ws.on_open = on_open

    ws.run_forever(ping_timeout=5)

    # ws.run_forever()


if __name__ == "__main__":

    run_1()
